import numpy as np
import random
import matplotlib.pyplot as plt


class PSO:

    """
    pso算法回顾
    计算公式为
    p[i] = p[i-1]+V[i]
    v[i] = w*v[i-1]+c1*r1*gbest[i]+c2*r2*pbest[i]
    """

    #初始化参数
    def __init__(self):
        # 给定迭代次数
        self.Number = 200

        # 初始化参数
        self.w = 0.6
        self.c1 = self.c2 = 2.0
        self.xMax = 10
        self.xMin = -10
        self.vMax = 5
        self.vMin = -5

        # 设置记录值
        self.maxSize = 50  # 粒子数
        self.dim = 3  # 粒子维数

        #画图
        plt.ion()

        self.xPos = np.empty([self.maxSize, self.dim], dtype="float64")
        self.xVel = np.empty([self.maxSize, self.dim], dtype="float64")
        self.fitNess = np.zeros([self.maxSize])

        self.pBest = np.empty([self.maxSize, self.dim], dtype="float64")
        self.gBest = np.empty([self.dim], dtype="float64")

    # 计算适应值
    def getResult(self,fitNess):
        result = 0
        for i in fitNess:
            result += i ** 2
        return result

    # 初始化
    def Create(self):

        # 设置粒子初始位置和初始速度
        for i in range(self.maxSize):
            for j in range(self.dim):
                self.xPos[i][j] = self.xMin + random.random() * (self.xMax - self.xMin)
                self.xVel[i][j] = self.vMin + random.random() * (self.vMax - self.vMin)

        # 初始化fitNess和pbest
        for i in range(self.maxSize):
            self.fitNess[i] = self.getResult(self.xPos[i])

        self.pBest = self.xPos

        # 初始化gbest
        bestFit = self.fitNess[0]
        self.gBest = self.xPos[0]
        for i in range(1, self.maxSize):
            if (bestFit > self.fitNess[i]):
                bestFit = self.fitNess[i]
                self.gBest = self.xPos[i]

        # 输出初始化信息
        print("初始化完成..........")
        print("0---->" + str(self.getResult(self.gBest)) + "->>>>", end="")
        for i in range(self.dim):
            print(str(self.gBest[i]) + "\t", end="")
        print(">>>>>")

    #更新状态
    def update(self):

        for j in range(self.Number):
            for i in range(self.maxSize):
                #位置和速度更新
                self.xVel[i] = self.xVel[i]*self.w+self.c1*random.random()*(self.gBest-self.xPos[i])+self.c2*random.random()*(self.pBest[i]-self.xPos[i])
                self.xPos[i] = self.xPos[i]+self.xVel[i]

                #计算fitNess,是否更新Pbest
                if self.getResult(self.xPos[i])<self.fitNess[i]:
                    self.pBest[i] = self.xPos[i]
                self.fitNess[i] = self.getResult(self.xPos[i])

                #判断是否更新gBest
                if self.fitNess[i]<self.getResult(self.gBest):
                    self.gBest = self.xPos[i]

            #输出信息
            print(str(j+1)+"---->" + str(self.getResult(self.gBest)) + "->>>>", end="")
            for k in range(self.dim):
                print(str(self.gBest[k]) + "\t", end="")
            print(">>>>>")

            self.show()
        plt.ioff()
        plt.show()
        plt.cla()

    #画图
    def show(self):
        # x = self.xPos.T[0]
        # y = self.xPos.T[1]
        # plt.figure(num="haha", figsize=(8, 7), facecolor="yellow", edgecolor="red", frameon=True)
        # plt.scatter(x, y)
        # # 设置轴的范围
        # plt.xlim(-12, 12)
        # plt.ylim(-10, 10)
        #
        # # 设置轴的含义
        # plt.xlabel("X", fontproperties="SimHei")
        # plt.xlabel("Y", fontproperties="SimHei")
        #
        # plt.show()

        x = self.xPos.T[0]
        y = self.xPos.T[1]

        plt.cla()
        plt.xlim(-10, 10)
        plt.ylim(-10,10)
        if 'sca' in globals():
            sca.remove()
        sca = plt.scatter(x,y, s=200, lw=0, c='red', alpha=0.5)
        plt.pause(0.05)

        # plt.show()

        # self.ax.cla()
        # self.ax.scatter(x,y)


if __name__ == "__main__":
    pso = PSO()
    pso.Create()
    pso.update()